"""知识库和文档管理接口"""
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
from typing import List, Optional
import os
import shutil
from pathlib import Path
import asyncio

from app.database import get_db
from app import models, schemas
from app.rag_service import rag_service
from app.config import settings
import logging

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/knowledge", tags=["知识库"])

# 文件上传目录
UPLOAD_DIR = Path("uploads")
UPLOAD_DIR.mkdir(exist_ok=True)


@router.post("/knowledge-bases", response_model=schemas.Response)
async def create_knowledge_base(
    name: str = Form(...),
    description: Optional[str] = Form(None),
    embedding_model: str = Form("bge-large-zh-v1.5"),
    vector_store: str = Form("faiss"),
    chunk_size: int = Form(500),
    chunk_overlap: int = Form(50),
    files: List[UploadFile] = File(...),
    db: Session = Depends(get_db)
):
    """创建知识库并上传文档"""
    try:
        # 1. 创建知识库
        kb = models.KnowledgeBase(
            name=name,
            description=description,
            embedding_model=embedding_model,
            vector_store=vector_store,
            config={
                "chunk_size": chunk_size,
                "chunk_overlap": chunk_overlap,
                "top_k": 3,
                "threshold": 0.5
            }
        )
        db.add(kb)
        db.commit()
        db.refresh(kb)
        
        logger.info(f"创建知识库: {kb.name} (ID: {kb.id})")
        
        # 2. 创建知识库专属目录
        kb_dir = UPLOAD_DIR / f"kb_{kb.id}"
        kb_dir.mkdir(exist_ok=True)
        
        # 3. 保存上传的文件
        uploaded_files = []
        for file in files:
            try:
                # 保存文件
                file_path = kb_dir / file.filename
                with open(file_path, "wb") as buffer:
                    shutil.copyfileobj(file.file, buffer)
                
                # 创建文档记录
                document = models.Document(
                    kb_id=kb.id,
                    title=file.filename,
                    file_path=str(file_path),
                    file_type=Path(file.filename).suffix.lstrip('.'),
                    file_size=os.path.getsize(file_path),
                    status="pending"
                )
                db.add(document)
                db.commit()
                db.refresh(document)
                
                uploaded_files.append({
                    "id": document.id,
                    "filename": file.filename,
                    "status": "pending"
                })
                
                logger.info(f"文件上传成功: {file.filename} (文档ID: {document.id})")
                
            except Exception as e:
                logger.error(f"文件上传失败 {file.filename}: {e}")
                continue
        
        # 4. 异步处理文档（向量化）
        # 在后台任务中处理，避免阻塞请求
        asyncio.create_task(process_documents_async(db, kb.id, uploaded_files))
        
        return schemas.Response(
            message=f"知识库创建成功，共上传 {len(uploaded_files)} 个文件，正在后台处理中",
            data={
                "kb_id": kb.id,
                "kb_name": kb.name,
                "files": uploaded_files
            }
        )
    
    except Exception as e:
        logger.error(f"创建知识库失败: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


async def process_documents_async(db: Session, kb_id: int, files: List[dict]):
    """异步处理文档"""
    for file_info in files:
        try:
            # 创建新的数据库会话
            from app.database import SessionLocal
            db_new = SessionLocal()
            
            await rag_service.process_document(
                db=db_new,
                document_id=file_info['id']
            )
            
            db_new.close()
            logger.info(f"文档处理完成: {file_info['filename']}")
        
        except Exception as e:
            logger.error(f"文档处理失败 {file_info['filename']}: {e}")
            continue


@router.get("/knowledge-bases", response_model=schemas.Response)
async def list_knowledge_bases(
    db: Session = Depends(get_db)
):
    """获取知识库列表"""
    kbs = db.query(models.KnowledgeBase).all()
    
    result = []
    for kb in kbs:
        # 统计文档数量和状态
        total_docs = db.query(models.Document).filter(
            models.Document.kb_id == kb.id
        ).count()
        
        completed_docs = db.query(models.Document).filter(
            models.Document.kb_id == kb.id,
            models.Document.status == "completed"
        ).count()
        
        result.append({
            "id": kb.id,
            "name": kb.name,
            "description": kb.description,
            "embedding_model": kb.embedding_model,
            "vector_store": kb.vector_store,
            "total_docs": total_docs,
            "completed_docs": completed_docs,
            "created_at": kb.created_at
        })
    
    return schemas.Response(data=result)


@router.get("/knowledge-bases/{kb_id}", response_model=schemas.Response)
async def get_knowledge_base(
    kb_id: int,
    db: Session = Depends(get_db)
):
    """获取知识库详情"""
    kb = db.query(models.KnowledgeBase).filter(
        models.KnowledgeBase.id == kb_id
    ).first()
    
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    
    # 获取文档列表
    documents = db.query(models.Document).filter(
        models.Document.kb_id == kb_id
    ).all()
    
    docs_data = []
    for doc in documents:
        docs_data.append({
            "id": doc.id,
            "title": doc.title,
            "file_type": doc.file_type,
            "file_size": doc.file_size,
            "status": doc.status,
            "error_msg": doc.error_msg,
            "created_at": doc.created_at
        })
    
    return schemas.Response(data={
        "id": kb.id,
        "name": kb.name,
        "description": kb.description,
        "embedding_model": kb.embedding_model,
        "vector_store": kb.vector_store,
        "config": kb.config,
        "documents": docs_data,
        "created_at": kb.created_at
    })


@router.post("/knowledge-bases/{kb_id}/documents", response_model=schemas.Response)
async def add_documents_to_kb(
    kb_id: int,
    files: List[UploadFile] = File(...),
    db: Session = Depends(get_db)
):
    """向已有知识库添加文档"""
    # 检查知识库是否存在
    kb = db.query(models.KnowledgeBase).filter(
        models.KnowledgeBase.id == kb_id
    ).first()
    
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    
    # 知识库目录
    kb_dir = UPLOAD_DIR / f"kb_{kb_id}"
    kb_dir.mkdir(exist_ok=True)
    
    # 保存文件
    uploaded_files = []
    for file in files:
        try:
            # 保存文件
            file_path = kb_dir / file.filename
            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
            
            # 创建文档记录
            document = models.Document(
                kb_id=kb_id,
                title=file.filename,
                file_path=str(file_path),
                file_type=Path(file.filename).suffix.lstrip('.'),
                file_size=os.path.getsize(file_path),
                status="pending"
            )
            db.add(document)
            db.commit()
            db.refresh(document)
            
            uploaded_files.append({
                "id": document.id,
                "filename": file.filename,
                "status": "pending"
            })
            
            logger.info(f"文件添加成功: {file.filename}")
            
        except Exception as e:
            logger.error(f"文件添加失败 {file.filename}: {e}")
            continue
    
    # 异步处理文档
    asyncio.create_task(process_documents_async(db, kb_id, uploaded_files))
    
    return schemas.Response(
        message=f"成功添加 {len(uploaded_files)} 个文件，正在后台处理中",
        data={"files": uploaded_files}
    )


@router.delete("/documents/{doc_id}", response_model=schemas.Response)
async def delete_document(
    doc_id: int,
    db: Session = Depends(get_db)
):
    """删除文档"""
    document = db.query(models.Document).filter(
        models.Document.id == doc_id
    ).first()
    
    if not document:
        raise HTTPException(status_code=404, detail="文档不存在")
    
    # 删除文件
    if document.file_path and os.path.exists(document.file_path):
        try:
            os.remove(document.file_path)
        except Exception as e:
            logger.error(f"删除文件失败: {e}")
    
    # 删除数据库记录（级联删除 chunks）
    db.delete(document)
    db.commit()
    
    return schemas.Response(message="文档删除成功")


@router.post("/search", response_model=schemas.Response)
async def search_knowledge_base(
    kb_id: int = Form(...),
    query: str = Form(...),
    top_k: int = Form(3),
    db: Session = Depends(get_db)
):
    """在知识库中搜索"""
    results = rag_service.search_similar_chunks(
        db=db,
        kb_id=kb_id,
        query=query,
        top_k=top_k
    )
    
    return schemas.Response(data=results)

